from concurrent import futures
import sys
import grpc

from app.rag.internal.logic.rag_graph import RagGraph
from app.rag.pb import rag_pb2, rag_pb2_grpc
from app.rag.internal.svc import svc
from common.utils import decorator, tracing

class RagServer(rag_pb2_grpc.RagServicer):
    def __init__(self, svcCtx: svc.ServiceContext):
        super().__init__()
        self.svcCtx = svcCtx
        self.tracing_util = tracing.TracingUtil()
        
        self._rag_graph = RagGraph(self.svcCtx)
    
    @decorator.my_grpc_handler(rag_pb2.AiReply)
    def ChatStream(self, request, context):
        # 创建带trace信息的metadata
        metadata = {}
        self.tracing_util.inject_trace_metadata(metadata)

        BUFFER_SIZE = 16
        buffer = []

        for resp in self._rag_graph.chat_stream(
            request.content, request.assistantId, request.dialogId, request.userId, metadata):
            buffer.append(resp)
            if len(buffer) >= BUFFER_SIZE:
                yield rag_pb2.AiReply(chunk="".join(buffer))
                buffer.clear()

        # 发送剩余未满缓冲区的数据
        if buffer:  
            yield rag_pb2.AiReply(chunk="".join(buffer))
            
    @decorator.my_grpc_handler(rag_pb2.AiReply)
    def Ping(self, request, context):
        return rag_pb2.Response(pong="1")

def serve(svcCtx: svc.ServiceContext):
    server = grpc.server(futures.ThreadPoolExecutor(max_workers=svcCtx.config.max_workers))
    rag_pb2_grpc.add_RagServicer_to_server(RagServer(svcCtx), server)
    server.add_insecure_port(svcCtx.config.listen_on)
    server.start()
    print(f"Starting rpc server 【RagServer】 at {svcCtx.config.listen_on}...\n", file=sys.stderr)
    server.wait_for_termination()